import chromadb
from chromadb.api.types import Documents, EmbeddingFunction, Embeddings
import openai
from sentence_transformers import SentenceTransformer
from typing import List


class SelfEmbeddingFunction(EmbeddingFunction):
    def __init__(self, embedding_func):
        self.embedding_func = embedding_func
    def __call__(self, texts: Documents) -> Embeddings:
        return self.embedding_func(texts)

class ChromaClient:
    def __init__(self, config: dict):
        if config['type'] == 'local':
            self.client = chromadb.PersistentClient(path=config['path'])
        else:
            self.client = chromadb.HttpClient(host=config['host'],port=config['port'])
        
        
        if config['embedding_type'] == 'local':
            self.model = SentenceTransformer(config['embedding_model'])
            def embedding_func(texts: List[str]):
                return self.model.encode(texts)
            
        else:

            self.openai_client = openai.OpenAI(base_url=config['base_url'], api_key=config['api_key'])
            def embedding_func(texts: List[str]):
                return [i.embedding for i in self.openai_client.embeddings.create(input=texts, model=config['embedding_model']).data]
            
        self.embedding_func = SelfEmbeddingFunction(embedding_func)

        self.config = config

    def create_collection(self, name: str):
        collection = self.client.get_or_create_collection(
            name=name, 
            embedding_function=self.embedding_func,
            metadata={
                "embedding_model": self.config['embedding_model']
            }
        )
        return collection
    
    def get_collection(self, name: str):
        collection = self.client.get_collection(name=name)
        return collection

    def delete_collection(self, name: str):
        self.client.delete_collection(name=name)

    def list_collections(self):
        collections = self.client.list_collections()
        for col in collections:
            print(col)
        return collections
    
    # 文档操作
    def add(self, collection_name: str, documents: List[str], metadatas: List[dict] = None):
        if metadatas is None:
            metadatas = [{} for _ in documents]
        collection = self.client.get_collection(name=collection_name)
        count = collection.count()
        ids = [str(count + i) for i in range(len(documents))]
        collection.add(documents=documents, metadatas=metadatas, ids=ids)

    def query(self, collection_name: str, query_texts: List[str], n_result: int = 5):
        collection = self.client.get_collection(name=collection_name)
        results = collection.query(query_texts=query_texts, n_results=n_result)
        return results

    def update(self, collection_name: str, document: str, metadata:dict, ids: str):
        collection = self.client.get_collection(name=collection_name)
        collection.update(documents=document, metadatas=metadata, ids=ids)
    
    def delete(self, collection_name: str, ids: List[str]):
        collection = self.client.get_collection(name=collection_name)
        collection.delete(ids=ids)

    def get(self, collection_name: str, ids: List[str]=None, where: dict = None, where_document: dict = None, limit: int=100):
        collection = self.client.get_collection(name=collection_name)
        if ids is None and where is None and where_document is None:
            return collection.get()
        else:
            return collection.get(ids=ids, where=where, where_document=where_document, limit=limit)
if __name__ == '__main__':
    config = {
        'type': 'remote',
        'host':'192.168.30.66',
        'port':8888,
        'embedding_type': 'openai',
        'embedding_model': 'BAAI/bge-large-zh-v1.5',
        'base_url':"https://api.siliconflow.cn/v1",
        'api_key':"sk-tkilzwjmiqcodtsovulcrkdmucgiqcabetzkbrglktvdtguo"
    }
    client = ChromaClient(config=config)
    print(client.list_collections())
    client.create_collection(name='aaaa')
    print(client.list_collections())
    documents = ['我爱我的祖国','我是中国人','我爱北京天安门']
    metadatas = [{'country':'China'},{'country':'China'},{'country':'China'}]
    client.add(collection_name='aaaa',documents=documents,metadatas=metadatas)
    print(client.query(collection_name='aaaa',query_texts=['我和我的祖国'],n_result=5))
    client.update(collection_name='aaaa',document='我很爱我的祖国',metadata={'country':'China'},ids='0')
    print(client.query(collection_name='aaaa',query_texts=['我很爱我的祖国'],n_result=5))
    client.delete(collection_name='aaaa',ids=['0','1','2'])
    client.delete_collection(name='aaaa')
    print(client.list_collections())
